/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia;



import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.unified.vendor.guava.com.google.common.net.HostAndPort;

import com.bff.gaia.api.common.ExecutionConfig;

import com.bff.gaia.api.common.JobExecutionResult;

import com.bff.gaia.api.java.CollectionEnvironment;

import com.bff.gaia.api.java.ExecutionEnvironment;

import com.bff.gaia.client.program.ClusterClient;

import com.bff.gaia.client.program.JobWithJars;

import com.bff.gaia.client.program.ProgramInvocationException;

import com.bff.gaia.client.program.rest.RestClusterClient;

import com.bff.gaia.configuration.*;

import com.bff.gaia.runtime.jobgraph.SavepointRestoreSettings;

import com.bff.gaia.runtime.state.StateBackend;

import com.bff.gaia.streaming.api.TimeCharacteristic;

import com.bff.gaia.streaming.api.environment.CheckpointConfig.ExternalizedCheckpointCleanup;

import com.bff.gaia.streaming.api.environment.RemoteStreamEnvironment;

import com.bff.gaia.streaming.api.environment.StreamExecutionEnvironment;

import com.bff.gaia.streaming.api.graph.StreamGraph;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import javax.annotation.Nullable;

import java.net.URL;

import java.util.Collections;

import java.util.List;



import static com.bff.gaia.streaming.api.environment.StreamExecutionEnvironment.getDefaultLocalParallelism;



/** Utilities for Gaia execution environments. */

public class GaiaExecutionEnvironments {



  private static final Logger LOG = LoggerFactory.getLogger(GaiaExecutionEnvironments.class);



  /**

   * If the submitted job is a batch processing job, this method creates the adequate Gaia {@link

   * com.bff.gaia.api.java.ExecutionEnvironment} depending on the user-specified options.

   */

  public static ExecutionEnvironment createBatchExecutionEnvironment(

	  GaiaPipelineOptions options, List<String> filesToStage) {

    return createBatchExecutionEnvironment(options, filesToStage, null);

  }



  static ExecutionEnvironment createBatchExecutionEnvironment(

	  GaiaPipelineOptions options, List<String> filesToStage, @Nullable String confDir) {



    LOG.info("Creating a Batch Execution Environment.");



    String masterUrl = options.getGaiaMaster();

    Configuration gaiaConfiguration = getGaiaConfiguration(confDir);

    ExecutionEnvironment gaiaBatchEnv;



    // depending on the master, create the right environment.

    if ("[local]".equals(masterUrl)) {

      gaiaBatchEnv = ExecutionEnvironment.createLocalEnvironment(gaiaConfiguration);

    } else if ("[collection]".equals(masterUrl)) {

      gaiaBatchEnv = new CollectionEnvironment();

    } else if ("[auto]".equals(masterUrl)) {

      gaiaBatchEnv = ExecutionEnvironment.getExecutionEnvironment();

    } else {

      int defaultPort = gaiaConfiguration.getInteger(RestOptions.PORT);

      HostAndPort hostAndPort = HostAndPort.fromString(masterUrl).withDefaultPort(defaultPort);

      gaiaConfiguration.setInteger(RestOptions.PORT, hostAndPort.getPort());

      gaiaBatchEnv =

          ExecutionEnvironment.createRemoteEnvironment(

              hostAndPort.getHost(),

              hostAndPort.getPort(),

              gaiaConfiguration,

              filesToStage.toArray(new String[filesToStage.size()]));

      LOG.info("Using Gaia Master URL {}:{}.", hostAndPort.getHost(), hostAndPort.getPort());

    }



    // Set the execution more for data exchange.

    gaiaBatchEnv.getConfig().setExecutionMode(options.getExecutionModeForBatch());



    // set the correct parallelism.

    if (options.getParallelism() != -1 && !(gaiaBatchEnv instanceof CollectionEnvironment)) {

      gaiaBatchEnv.setParallelism(options.getParallelism());

    }

    // Set the correct parallelism, required by UnboundedSourceWrapper to generate consistent

    // splits.

    final int parallelism;

    if (gaiaBatchEnv instanceof CollectionEnvironment) {

      parallelism = 1;

    } else {

      parallelism =

          determineParallelism(

              options.getParallelism(), gaiaBatchEnv.getParallelism(), gaiaConfiguration);

    }



    gaiaBatchEnv.setParallelism(parallelism);

    // set parallelism in the options (required by some execution code)

    options.setParallelism(parallelism);



    if (options.getObjectReuse()) {

      gaiaBatchEnv.getConfig().enableObjectReuse();

    } else {

      gaiaBatchEnv.getConfig().disableObjectReuse();

    }



    applyLatencyTrackingInterval(gaiaBatchEnv.getConfig(), options);



    return gaiaBatchEnv;

  }



  /**

   * If the submitted job is a stream processing job, this method creates the adequate Gaia {@link

   * com.bff.gaia.streaming.api.environment.StreamExecutionEnvironment} depending on the

   * user-specified options.

   */

  public static StreamExecutionEnvironment createStreamExecutionEnvironment(

	  GaiaPipelineOptions options, List<String> filesToStage) {

    return createStreamExecutionEnvironment(options, filesToStage, null);

  }



  @VisibleForTesting

  static StreamExecutionEnvironment createStreamExecutionEnvironment(

	  GaiaPipelineOptions options, List<String> filesToStage, @Nullable String confDir) {



    LOG.info("Creating a Streaming Environment.");



    String masterUrl = options.getGaiaMaster();

    Configuration gaiaConfiguration = getGaiaConfiguration(confDir);

    final StreamExecutionEnvironment gaiaStreamEnv;



    // depending on the master, create the right environment.

    if ("[local]".equals(masterUrl)) {

      gaiaStreamEnv =

          StreamExecutionEnvironment.createLocalEnvironment(

              getDefaultLocalParallelism(), gaiaConfiguration);

    } else if ("[auto]".equals(masterUrl)) {

      gaiaStreamEnv = StreamExecutionEnvironment.getExecutionEnvironment();

    } else {

      int defaultPort = gaiaConfiguration.getInteger(RestOptions.PORT);

      HostAndPort hostAndPort = HostAndPort.fromString(masterUrl).withDefaultPort(defaultPort);

      gaiaConfiguration.setInteger(RestOptions.PORT, hostAndPort.getPort());

      final SavepointRestoreSettings savepointRestoreSettings;

      if (options.getSavepointPath() != null) {

        savepointRestoreSettings =

            SavepointRestoreSettings.forPath(

                options.getSavepointPath(), options.getAllowNonRestoredState());

      } else {

        savepointRestoreSettings = SavepointRestoreSettings.none();

      }

      gaiaStreamEnv =

          new UnifiedGaiaRemoteStreamEnvironment(

              hostAndPort.getHost(),

              hostAndPort.getPort(),

              gaiaConfiguration,

              savepointRestoreSettings,

              filesToStage.toArray(new String[filesToStage.size()]));

      LOG.info("Using Gaia Master URL {}:{}.", hostAndPort.getHost(), hostAndPort.getPort());

    }



    // Set the parallelism, required by UnboundedSourceWrapper to generate consistent splits.

    final int parallelism =

        determineParallelism(

            options.getParallelism(), gaiaStreamEnv.getParallelism(), gaiaConfiguration);

    gaiaStreamEnv.setParallelism(parallelism);

    if (options.getMaxParallelism() > 0) {

      gaiaStreamEnv.setMaxParallelism(options.getMaxParallelism());

    }

    // set parallelism in the options (required by some execution code)

    options.setParallelism(parallelism);



    if (options.getObjectReuse()) {

      gaiaStreamEnv.getConfig().enableObjectReuse();

    } else {

      gaiaStreamEnv.getConfig().disableObjectReuse();

    }



    // default to event time

    gaiaStreamEnv.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);



    // for the following 2 parameters, a value of -1 means that Gaia will use

    // the default values as specified in the configuration.

    int numRetries = options.getNumberOfExecutionRetries();

    if (numRetries != -1) {

      gaiaStreamEnv.setNumberOfExecutionRetries(numRetries);

    }

    long retryDelay = options.getExecutionRetryDelay();

    if (retryDelay != -1) {

      gaiaStreamEnv.getConfig().setExecutionRetryDelay(retryDelay);

    }



    // A value of -1 corresponds to disabled checkpointing (see CheckpointConfig in Gaia).

    // If the value is not -1, then the validity checks are applied.

    // By default, checkpointing is disabled.

    long checkpointInterval = options.getCheckpointingInterval();

    if (checkpointInterval != -1) {

      if (checkpointInterval < 1) {

        throw new IllegalArgumentException("The checkpoint interval must be positive");

      }

      gaiaStreamEnv.enableCheckpointing(checkpointInterval, options.getCheckpointingMode());

      if (options.getCheckpointTimeoutMillis() != -1) {

        gaiaStreamEnv

            .getCheckpointConfig()

            .setCheckpointTimeout(options.getCheckpointTimeoutMillis());

      }

      boolean externalizedCheckpoint = options.isExternalizedCheckpointsEnabled();

      boolean retainOnCancellation = options.getRetainExternalizedCheckpointsOnCancellation();

      if (externalizedCheckpoint) {

        gaiaStreamEnv

            .getCheckpointConfig()

            .enableExternalizedCheckpoints(

                retainOnCancellation

                    ? ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION

                    : ExternalizedCheckpointCleanup.DELETE_ON_CANCELLATION);

      }



      long minPauseBetweenCheckpoints = options.getMinPauseBetweenCheckpoints();

      if (minPauseBetweenCheckpoints != -1) {

        gaiaStreamEnv

            .getCheckpointConfig()

            .setMinPauseBetweenCheckpoints(minPauseBetweenCheckpoints);

      }

      boolean failOnCheckpointingErrors = options.getFailOnCheckpointingErrors();

      gaiaStreamEnv.getCheckpointConfig().setFailOnCheckpointingErrors(failOnCheckpointingErrors);

    }



    applyLatencyTrackingInterval(gaiaStreamEnv.getConfig(), options);



    if (options.getAutoWatermarkInterval() != null) {

      gaiaStreamEnv.getConfig().setAutoWatermarkInterval(options.getAutoWatermarkInterval());

    }



    // State backend

    final StateBackend stateBackend = options.getStateBackend();

    if (stateBackend != null) {

      gaiaStreamEnv.setStateBackend(stateBackend);

    }



    return gaiaStreamEnv;

  }



  private static int determineParallelism(

      final int pipelineOptionsParallelism,

      final int envParallelism,

      final Configuration configuration) {

    if (pipelineOptionsParallelism > 0) {

      return pipelineOptionsParallelism;

    }

    if (envParallelism > 0) {

      // If the user supplies a parallelism on the command-line, this is set on the execution

      // environment during creation

      return envParallelism;

    }



    final int gaiaConfigParallelism =

        configuration.getInteger(CoreOptions.DEFAULT_PARALLELISM.key(), -1);

    if (gaiaConfigParallelism > 0) {

      return gaiaConfigParallelism;

    }

    LOG.warn(

        "No default parallelism could be found. Defaulting to parallelism 1. "

            + "Please set an explicit parallelism with --parallelism");

    return 1;

  }



  private static Configuration getGaiaConfiguration(@Nullable String gaiaConfDir) {

    return gaiaConfDir == null

        ? GlobalConfiguration.loadConfiguration()

        : GlobalConfiguration.loadConfiguration(gaiaConfDir);

  }



  private static void applyLatencyTrackingInterval(

      ExecutionConfig config, GaiaPipelineOptions options) {

    long latencyTrackingInterval = options.getLatencyTrackingInterval();

    config.setLatencyTrackingInterval(latencyTrackingInterval);

  }



  /**

   * Remote stream environment that supports job execution with restore from savepoint.

   *

   * <p>This class can be removed once Gaia provides this functionality.

   *

   * <p>TODO: https://issues.apache.org/jira/browse/BEAM-5396

   */

  private static class UnifiedGaiaRemoteStreamEnvironment extends RemoteStreamEnvironment {

    private final SavepointRestoreSettings restoreSettings;



    public UnifiedGaiaRemoteStreamEnvironment(

        String host,

        int port,

        Configuration clientConfiguration,

        SavepointRestoreSettings restoreSettings,

        String... jarFiles) {

      super(host, port, clientConfiguration, jarFiles, null);

      this.restoreSettings = restoreSettings;

    }



    // copied from RemoteStreamEnvironment and augmented to pass savepoint restore settings

    @Override

    protected JobExecutionResult executeRemotely(StreamGraph streamGraph, List<URL> jarFiles)

        throws ProgramInvocationException {



      List<URL> globalClasspaths = Collections.emptyList();

      String host = super.getHost();

      int port = super.getPort();



      if (LOG.isInfoEnabled()) {

        LOG.info("Running remotely at {}:{}", host, port);

      }



      ClassLoader usercodeClassLoader =

          JobWithJars.buildUserCodeClassLoader(

              jarFiles, globalClasspaths, getClass().getClassLoader());



      Configuration configuration = new Configuration();

      configuration.addAll(super.getClientConfiguration());



      configuration.setString(JobManagerOptions.ADDRESS, host);

      configuration.setInteger(JobManagerOptions.PORT, port);



      configuration.setInteger(RestOptions.PORT, port);



      final ClusterClient<?> client;

      try {

        client = new RestClusterClient<>(configuration, "RemoteStreamEnvironment");

      } catch (Exception e) {

        throw new ProgramInvocationException(

            "Cannot establish connection to JobManager: " + e.getMessage(), e);

      }



      client.setPrintStatusDuringExecution(getConfig().isSysoutLoggingEnabled());



      try {

        return client

            .run(streamGraph, jarFiles, globalClasspaths, usercodeClassLoader, restoreSettings)

            .getJobExecutionResult();

      } catch (ProgramInvocationException e) {

        throw e;

      } catch (Exception e) {

        String term = e.getMessage() == null ? "." : (": " + e.getMessage());

        throw new ProgramInvocationException("The program execution failed" + term, e);

      } finally {

        try {

          client.shutdown();

        } catch (Exception e) {

          LOG.warn("Could not properly shut down the cluster client.", e);

        }

      }

    }

  }

}